# coding:utf-8
from associator.associator import Associator
from efficient_apriori import apriori
import pandas as pd

'''
Apriori算法
'''


class Apriori(Associator):

    def __init__(self):
        Associator.__init__(self)
        self.algorithm_name = "Apriori算法"
        self.ipynb_template_name = "apriori-template.ipynb"
    
    def implent(self):
        Associator.implent(self)
        self.itemsets, self.rules = apriori(self.transactions, min_support=self.min_support, min_confidence=self.min_confidence)
    
    def saveToExcle(self):
        # 保存结果excel文件
        data_array = []
        for rule in self.rules:
            data_array.append(["+".join(rule.lhs), "+".join(rule.rhs), rule.support, rule.confidence, len(rule.lhs), len(rule.rhs)])
        pd.DataFrame(data=data_array, columns=["前件", "后件", "支持度", "可信度", "前件项数", "后件项数"]).to_excel(self.result_excel_path)

